import os
import uuid
import json
import time
import jwt
import requests
import cv2
from typing import Optional, Tuple, Dict, Any

import urllib3
from dotenv import load_dotenv
from openai import OpenAI

from utils import sr_s3, status_codes
from utils.sr_utils import image_path_to_data_url, content_moderations, video_compress_image_size_path

StatusCodes = status_codes.StatusCodes


class viduGenerator_Vid:
    """视频生成工具类，封装KlingAI视频生成API调用逻辑，支持外部调用"""

    def __init__(self):
        """初始化配置与环境变量"""
        # 加载环境变量
        self.verify_ssl = True
        load_dotenv(verbose=True, override=True)
        self.app_env = os.getenv('APP_ENV', 'prod')
        self.url = "https://api.vidu.com"
        self.viduKey = os.getenv('viduKey')
        self.KELING_SK = os.getenv('KELING_SK')
        self.NSFW_CHECK_ENABLED = os.getenv('NSFW_CHECK_ENABLED')
        self.MODELS = {
            1: "viduq1",
            2: "vidu2.0",
            3: "vidu1.5",
            4: "viduq1-classic",
        }
        # 初始化鉴权头（JWT动态生成，30分钟有效期）
        self.headers = {
            "Authorization": f"Token {self.viduKey}",
            "Content-Type": "application/json"
        }
        self.headers2 = {
            "Authorization": f"Token {self.viduKey}",
        }
        # 临时文件存储目录（可外部配置）
        self.temp_image_dir = "./upload"
        self.temp_video_dir = "./download"
        os.makedirs(self.temp_image_dir, exist_ok=True)
        os.makedirs(self.temp_video_dir, exist_ok=True)

    # --------------------------
    # 内部工具方法：鉴权与配置
    # --------------------------

    # def _encode_jwt_token(self):
    #     headers = {
    #         "alg": "HS256",
    #         "typ": "JWT"
    #     }
    #     payload = {
    #         "iss": "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BAp",
    #         "exp": int(time.time()) + 1800,  # 有效时间，此处示例代表当前时间+1800s(30min)
    #         "nbf": int(time.time()) - 5  # 开始生效的时间，此处示例代表当前时间-5秒
    #     }
    #     token = jwt.encode(payload, "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BA?             p", headers=headers)
    #     return token

    # --------------------------
    # 内部工具方法：文件处理
    # --------------------------
    def _delete_temp_files(self, *file_paths: str) -> None:
        """删除临时文件（忽略不存在的文件）"""
        for path in file_paths:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"清理临时文件失败 {path}: {str(e)}")

    def _download_and_compress_image(self, img_url: str, storage_service: str, bucket_name: str) -> Tuple[bool, str, int]:

        from utils import sr_s3
        """下载图片并压缩（返回是否成功、base64编码、是否压缩）"""

        if self.NSFW_CHECK_ENABLED == "true":
            image_url = sr_s3.get_url(img_url, serviceName=storage_service, bucketName=bucket_name)
            clientOpenAI = OpenAI()
            image_result = content_moderations(clientOpenAI, image_url)
            if image_result:
                return False, f"图片违规", 5
        try:
            # 生成临时保存路径
            name = str(uuid.uuid4())
            ext = os.path.splitext(img_url)[-1].lower() or '.jpg'
            save_path = os.path.join(self.temp_image_dir, f"{name}{ext}")
            # 调用S3工具下载图片（假设sr_s3为外部存储工具）
            success, err = sr_s3.download_file(
                img_url, save_path, serviceName=storage_service, bucketName=bucket_name
            )
            if not success:
                return False, f"下载失败: {err}", 2
            # 压缩图片（确保符合模型要求）
            success, new_path = video_compress_image_size_path(save_path, 10, (300, 6000), (300, 6000))
            if not success:
                return False, "图片压缩失败", 2
            base64_img = image_path_to_data_url(new_path)
            return True, base64_img, 1
        except Exception as e:
            return False, f"图片处理异常: {str(e)}", 2

    def _download_video(self, video_url: str, task_id: str) -> Tuple[bool, Dict[str, Any]]:
        """下载生成的视频到本地（返回是否成功、视频信息）"""
        try:
            filename = f"video_{task_id}.mp4"
            save_path = os.path.join(self.temp_video_dir, filename)

            # 流式下载视频
            response = requests.get(video_url, stream=True, timeout=30)
            response.raise_for_status()

            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB/块
                    if chunk:
                        f.write(chunk)

            # 获取视频信息（宽高、帧率等）
            cap = cv2.VideoCapture(save_path)
            video_info = {
                "save_path": save_path,
                "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "fps": cap.get(cv2.CAP_PROP_FPS),
                "duration": cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0
            }
            cap.release()
            return True, video_info
        except Exception as e:
            return False, {"error": f"视频下载失败: {str(e)}"}

    def _upload_video_to_s3(self, video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        """上传视频到S3（返回是否成功、S3路径）"""
        if not video_path or not os.path.exists(video_path):
            return False, "视频文件不存在"
        try:
            from utils import sr_s3
            s3_key = sr_s3.upload_login_or_visitor(
                video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
            )
            return True, s3_key
        except Exception as e:
            return False, f"上传S3失败: {str(e)}"

    # --------------------------
    # 核心方法：API调用与任务处理
    # --------------------------
    def _create_vidu_video_task(self, payload: dict, api_path: str) -> Tuple[bool, str, str]:
        """调用KlingAI API创建视频任务（返回是否成功、任务ID、错误信息）"""
        try:
            url = f"{self.url}{api_path}"
            print(url)
            response = requests.post(url, headers=self.headers, data=json.dumps(payload))
            response_data = response.json()
            print(response_data)
            if response_data.get("state") == 'created':
                task_id = response_data.get("task_id")
                if task_id:
                    return True, task_id, ""
                else:
                    return False, "", "API返回成功但无task_id"

            else:
                return False, response_data.get("state"), response_data.get("err_code"),
        except Exception as e:
            return False, "", f"创建任务失败: {str(e)}"

    def _poll_task_status(self, task_id: str, api_path: str, timeout: int = 600, interval: int = 5) -> Tuple[bool, str]:
        """轮询任务状态（返回是否成功、视频URL/错误信息）"""
        print("轮询开始")
        if not task_id:
            return False, "任务ID为空"
        url = f"https://api.vidu.com/ent/v2/tasks/{task_id}/creations"
        start_time = time.time()
        retries = 0
        max_retries = 5  # 最大重试次数
        # 创建带重试机制的Session
        session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            max_retries=max_retries,
            pool_connections=10,
            pool_maxsize=10
        )
        session.mount('https://', adapter)

        # 禁用不安全请求的警告
        if not self.verify_ssl:
            urllib3.disable_warnings()

        while time.time() - start_time < timeout:
            try:
                # 发送请求（增加超时设置和SSL验证控制）
                response = session.get(
                    url,
                    headers=self.headers2,
                    timeout=(5, 15),  # (连接超时, 读取超时)
                    verify=self.verify_ssl  # 控制是否验证SSL证书
                )
                print(response.json())
                response.raise_for_status()
                response_data = response.json()
                # print(f"轮询响应: {response_data}")
                # code = response_data.get("code")  # 401 = 未经授权
                #
                # reason = response_data.get("reason")  # 401 = 未经授权
                task_status = response_data.get("state")
                message = response_data.get("err_code", "无信息")
                if task_status == "success":
                    videos = response_data.get("creations", {})
                    if videos and "url" in videos[0]:
                        return True, videos[0]["url"]
                    else:
                        return False, "任务成功但无视频URL"
                elif task_status == "failed":
                    sr_s3.set_logging('vido任务失败:' + json.dumps(response_data))
                    return False, message
                else:
                    # 任务处理中，继续轮询
                    print(f"任务状态: {task_status}，等待{interval}秒后重试")
                    time.sleep(interval)
                    retries = 0  # 重置重试计数

            except requests.exceptions.SSLError as e:
                retries += 1
                if retries >= max_retries:
                    return False, f"SSL错误（重试{retries}次后）: {str(e)}"
                print(f"SSL连接错误（{retries}/{max_retries}）: {str(e)}，{interval}秒后重试...")
                time.sleep(interval * retries)  # 指数退避策略

            except requests.exceptions.RequestException as e:
                retries += 1
                if retries >= max_retries:
                    return False, f"网络请求错误（重试{retries}次后）: {str(e)}"
                print(f"网络请求错误（{retries}/{max_retries}）: {str(e)}，{interval}秒后重试...")
                time.sleep(interval * retries)  # 指数退避策略

            except json.JSONDecodeError:
                retries += 1
                if retries >= max_retries:
                    return False, f"响应格式错误（重试{retries}次后）: 无法解析JSON"
                print(f"响应格式错误（{retries}/{max_retries}）: 无法解析JSON，{interval}秒后重试...")
                time.sleep(interval * retries)

            except Exception as e:
                return False, f"轮询异常: {str(e)}"

        return False, f"任务超时（超过{timeout}秒）"

    # --------------------------
    # 外部调用方法：主入口
    # --------------------------
    def process_video_task_vidu(self, task_params: Dict[str, Any]):
        """
        处理视频生成任务（外部调用入口）

        参数:
            task_params: 任务参数字典，包含：
                - template : 模板名称
                - taskId: 任务ID
                - moduleName: 模型名称索引
                - prompt: 提示词
                - seed 随机种子 当默认不传或者传0时，
                - modelType: 模型类型（1:文生视频, 2:单图生视频, 3:首尾图生视频, 4:模板, 5:多图特效）
                - duration: 视频时长（默认5）
                - first_image: 首图URL（modelType=2/3/4/5时需要）
                - last_image: 尾图URL（modelType=3/5时需要）
                - storageService: 存储服务类型（如'r2'）
                - bucketName: 存储桶名称
                - bgm:是否为生成的视频添加背景音乐。默认：false，可选值 true 、false
                - storedPrefix: S3存储路径前缀
                - QueueResult: 结果队列名称
                - 其他可选参数：fps, resolution, ratio等

        返回:
            (是否成功, 结果字典) 结果字典包含's3_key'（视频S3路径）或'error'（错误信息）
        """
        # 初始化变量
        first_image_path = None
        last_image_path = None
        video_info = None

        try:
            # 1. 参数校验
            # required_params = ['modelType', 'storageService', 'bucketName', 'storedPrefix', 'QueueResult']
            # for param in required_params:
            #     if not task_params.get(param):
            #         return {"status": 2, "msg": f"缺少参数: {param}", 'code': 6003}

            model_type = task_params.get('modelType')
            modelName = task_params.get('moduleName')
            # 2. 构建API请求参数与路径
            api_path = ""
            payload = {"prompt": task_params.get('prompt'), "seed": task_params.get('seed', 0), 'duration': task_params.get('duration'), 'resolution': task_params.get('mResolution', '1080p'),
                       'movement_amplitude': task_params.get('motion_type', 'auto')}

            if payload['resolution'] == '':
                payload.pop('resolution', None)
            if task_params.get('bgm'):
                payload['bgm'] = True
            else:
                payload['bgm'] = False
            # 根据模型类型补充参数
            if model_type == 1:  # 文生视频
                api_path = "/ent/v2/text2video"
                payload['model'] = modelName
                payload["aspect_ratio"] = task_params.get('ratio', "16:9")
                payload['style'] = task_params.get('model_style', "general")
                if payload['aspect_ratio'] == '':
                    payload.pop('aspect_ratio', None)
            elif model_type == 2:  # 单图生视频/单图特效
                payload['model'] = modelName
                api_path = "/ent/v2/img2video"
                payload['resolution'] = task_params.get('mResolution', '1080p')
                first_image_url = task_params.get('first_image')
                if not first_image_url:
                    return {"status": 2, "msg": "modelType=2/4需提供first_image", 'code': 6003}
                # 下载并处理首图
                first_image = sr_s3.get_url(first_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                payload['images'] = [first_image]
            elif model_type == 3:  # 首尾图生视频
                api_path = "/ent/v2/start-end2video"
                payload['model'] = modelName
                first_image_url = task_params.get('first_image')
                last_image_url = task_params.get('last_image')
                if not (first_image_url and last_image_url):
                    return {"status": 2, "msg": "modelType=3/5需提供first_image和last_image", 'code': 6005}
                first_image = sr_s3.get_url(first_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                last_image = sr_s3.get_url(last_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                payload["images"] = [first_image, last_image]

            elif model_type == 4:  # 模板 去其他区分开 重写 payload
                payload = {"prompt": task_params.get('prompt'), "seed": task_params.get('seed', 0)}
                if task_params.get('bgm'):
                    payload['bgm'] = True
                else:
                    payload['bgm'] = False
                api_path = "/ent/v2/template2video"
                payload['template'] = task_params.get('template')
                if task_params.get('area'):
                    payload['area'] = task_params.get('area')
                if task_params.get('beast'):
                    payload['beast'] = task_params.get('beast')
                first_image_url = task_params.get('first_image')
                last_image_url = task_params.get('last_image')
                if task_params.get('template') == "relax_cut":
                    payload.pop("prompt")
                    payload["extra_params"] = {"object": task_params.get('prompt')}
                elif task_params.get('template') == "smooth_shift" or task_params.get('template') == "outfit_show":
                    if first_image_url is not None and last_image_url is not None:
                        first_img = sr_s3.get_url(first_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                        last_image = sr_s3.get_url(last_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                        payload['images'] = [first_img, last_image]
                elif task_params.get('template') == "beast_companion":
                    if first_image_url is not None:
                        first_img = sr_s3.get_url(first_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                        payload['images'] = [first_img]
                    payload['beast'] = task_params.get('beast', 'auto')
                else:
                    if first_image_url is not None:
                        first_img = sr_s3.get_url(first_image_url, serviceName=task_params['storageService'], bucketName=task_params['bucketName'])
                        payload['images'] = [first_img]

                if task_params.get('template') == "love_story":
                    del payload['template']
                    api_path = "/ent/v2/template-story"
                    payload['story'] = "love_story"
            elif model_type == 5:
                payload['model'] = modelName
                payload['images'] = []
                for i in task_params['images']:
                    success, first_img, status = self._download_and_compress_image(
                        i, task_params['storageService'], task_params['bucketName']
                    )
                    if not success:
                        return {"status": status, "msg": first_img, "code": 6005}
                    payload['images'].append(first_img)

            else:
                return {"status": 2, "msg": f"不支持的模型类型: {model_type}"}
            # 3. 创建视频任务
            print(payload)
            # exit()
            success, task_api_id, err_msg = self._create_vidu_video_task(payload, api_path)
            if not success:
                if err_msg == "AuditSubmitIllegal" and err_msg == "CreationPolicyViolation":
                    return {"status": 7, "msg": err_msg, "code": StatusCodes.INPUT_IMAGE_INVALID_MODEL}
                return {"status": 2, "msg": err_msg}
            # task_api_id = '845423452071407616'
            # print(f"任务{task_id}已提交至API，API任务ID: {task_api_id}")

            # 4. 轮询任务结果
            success, video_url = self._poll_task_status(
                task_api_id, api_path,
                timeout=task_params.get('timeout', 1800),
                interval=task_params.get('interval', 5)
            )

            print(video_url)
            if not success:
                if video_url == "AuditSubmitIllegal":
                    return {"status": 7, "msg": video_url, "code": StatusCodes.INPUT_IMAGE_INVALID_MODEL}
                elif video_url == "TaskPromptPolicyViolation":
                    return {"status": 7, "msg": video_url, "code": StatusCodes.INPUT_IMAGE_INVALID_MODEL}
                elif video_url == "CreationPolicyViolation":
                    return {"status": 7, "msg": video_url, "code": StatusCodes.OUTPUT_VIDEO_INVALID_MODEL}
                elif video_url == "ImageDownloadFailure":
                    return {"status": 2, "msg": video_url, "code": StatusCodes.IMAGE_PROCESS_ERROR_MODEL}
                return {"status": 2, "msg": video_url}
            # 5. 下载视频并上传到S3
            success, video_info = self._download_video(video_url, task_api_id)
            if not success:
                return {"status": 2, "msg": video_info.get("error"), 'code': 6001}
            print("下载视频地址", video_info["save_path"])
            # 6. 上传视频到S3
            success, s3_key = self._upload_video_to_s3(
                video_info["save_path"],
                task_params['storedPrefix'],
                task_params['storageService'],
                task_params['bucketName']
            )
            if not success:
                return {"status": 2, "msg": s3_key}

            # 7. 发送结果到队列（假设sr_sqs为队列工具）
            from utils import sr_sqs
            msg_body = {
                'status': 1,
                'width': video_info['width'],
                'height': video_info['height'],
                'local_path': video_info['save_path'],
                'url': s3_key,
                "model_mode": "std",
                'msg': "success",
                "code": ""
            }
            return msg_body

        except Exception as e:
            error_msg = f"任务处理异常: {str(e)}"

            return {"status": 2, "msg": error_msg, 'code': 6001}


# --------------------------
# 外部调用示例
# --------------------------
if __name__ == "__main__":
    # 初始化视频生成器
    video_generator = viduGenerator_Vid()
    # 示例任务参数
    # task_params = {'taskId': '16c33c4b6dfd9a220a53', 'prompt': None, 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754995108557.jpg', 'last_image': '', 'duration': '0', 'ratio': 'adaptive', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 4, 'moduleType': 24, 'mResolution': '720p', 'template': 'fisheye_human', 'p_t_b': 0, 'bgm': 1}
    # task_params = {'taskId': 'a4867698078b2ba29213', 'prompt': "Video content\\\\n The video shows a man taking off his shirt, revealing his muscular chest.\\\\n# Requirements\\\\n1. If the garment is a shirt, he would first unbutton it before removing it, revealing his toned muscles underneath.\\\\n2. Make sure that after the clothes are removed, they are thrown to one side of the frame.\\\\n3. Motion Level：Large.\\\\n4. The description of the 'Subject' should focus on the action of the person taking off the clothes, and then throwing the clothes aside while showing off their muscles.", 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1753962353536.png', 'last_image': '', 'duration': '0', 'ratio': 'adaptive', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 4, 'moduleType': 24, 'mResolution': '720P', 'template': 'muscling', 'bgm': 1}
    # task_params = {'taskId': '4d2a274dfd0fe53b61e2', 'prompt': 'Video content\\\\n The character in the image floats and flies like a superhero.\\\\n# Requirements\\\\n1.Camera Movement：track-up shot.\\\\n2.Motion Level：Large.', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754937103834.jpg', 'last_image': '', 'duration': '0', 'ratio': 'adaptive', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1,
    task_params = {
        'taskId' : 'fc96518d7411dded25f0',
    'prompt' : '屋里有人打开了窗户',
    'first_image' : 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1755763196906.jpg',
    'last_image' : '',
    'duration' : '4',
    'fps' : '25',
    'resolution' : '360p',
    'ratio' : '',
    'appId' : '1',
    'taskType' : 'pwai_video_itf',
    'QueueResult' : 'AILocalVideoResult',
    'storageService' : 'r2',
    'isLogin' : 1,
    'bucketName' : 'picwand',
    'storedPrefix' : 'video/image_to_video',
    'progress' : 1,
    'queueName' : 'AIHubVideoTasks_Test',
    'modelType' : 2,
    'moduleType' : 24,
    'moduleName' : 'vidu2.0',
    'mResolution' : '360p',
    'mode' : 'std',
    'negative_prompt' : '',
    'cfg_scale' : 0.5,
    'p_t_b' : 0,
    'n_t_b' : 0,
    'motion_type' : 'auto',
    'model_style' : 'general',
    'camera_movement' : '',
    }
    #                'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video',
    #                'progress': 1, 'queueName':
    #                'AIHubVideoTasks_Test',
    #                'modelType': 4, 'moduleType': 24,
    #                'mResolution': '720P', 'template': 'flying', 'bgm': 0}

    task_params = {
        'taskId': '1758764395144DrvOkmfDY64Xrh41',
        'prompt': '',
        'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1758784454467.jpg',
        'last_image': '',
        'duration': '4',
        'resolution': '720p',
        'ratio': 'adaptive',
        'appId': '1',
        'taskType': 'pwai_video_itf',
        'QueueResult': 'AILocalVideoResult',
        'storageService': 'r2',
        'isLogin': 1,
        'bucketName': 'picwand',
        'storedPrefix': 'video/image_to_video',
        'progress': 1,
        'queueName': 'AIHubVideoTasks_Test',
        'modelType': 4,
        'moduleType': 24,
        'mResolution': '360p',
        'template': 'love_story',
        'p_t_b': 0,
        'beast': '',
        'bgm': 1,
    }

    print(task_params)
    result = video_generator.process_video_task_vidu(task_params)
    print(result)
